#
import argparse
from typing import Dict
import os
from pathlib import Path
import json
import datetime
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as T
import torch.backends.cudnn as cudnn
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
import matplotlib.pyplot as plt
from torch.optim import Adam
from datasets import load_dataset
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
# Token Statistics Self Attention Model
from anns.tost.vis.datasets import build_dataset, build_transform
from anns.tost.vis.samplers import RASampler
from anns.tost.vis.engine import train_one_epoch, evaluate, predict
from anns.tost.vis.losses import DistillationLoss
from timm.data import Mixup
from timm.models import create_model
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.scheduler import create_scheduler
from timm.optim import create_optimizer
from timm.utils import NativeScaler, get_state_dict, ModelEma
import anns.tost.vis.tost
#
from apps.clip.positional_embedding import PositionalEmbedding
from apps.clip.attention_head import AttentionHead, MultiHeadAttention
from apps.clip.fashion_mnist import FashionMNIST
from apps.clip.clip_model import CLIP
from apps.clip.clip_config import ClipConfig as CC
from apps.clip.token import tokenizer
from anns.tost.vis import utils as TVU
from apps.clip.cifar100_text import Cifar100Text

class ClipApp(object):
    def __init__(self):
        self.name = 'apps.clip.clip_app.ClipApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'White Box Transformer(Tost) Clip')
        CC.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if params['task'] == 1:
            ClipApp.train_main(params=params)
        elif params['task'] == 2:
            ClipApp.test_main(params=params)
        elif params['task'] == 3:
            ClipApp.zero_shot(params=params)
        elif params['task'] == 4:
            # Tost训练过程
            ClipApp.tost_main()
        elif params['task'] == 5:
            ClipApp.prepare_cifar100_text(params=params)

    @staticmethod
    def prepare_cifar100_text(params:Dict = {}) -> None:
        mm_text_en, mm_text_zh = Cifar100Text.load_cifar100_texts()
        print(f'英文列表：')
        for k,v in mm_text_en.items():
            print(f'### {k}: {v};')
        for k,v in mm_text_zh.items():
            print(f'@@@ {k}: {v};')

    @staticmethod
    def train_main(params:Dict = {}) -> None:
        # 加载数据集
        train_set = FashionMNIST(train = True)
        test_set = FashionMNIST(train = False)
        train_loader = DataLoader(train_set, shuffle=True, batch_size=CC.batch_size)
        test_loader = DataLoader(test_set, shuffle=False, batch_size=CC.batch_size)
        print(train_set[0]["image"].shape)
        print(train_set[0]["caption"].shape)
        print(train_set[0]["mask"].shape)
        print(train_set[0]["caption"])
        print(train_set[0]["mask"])
        # 训练过程
        CC.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("Using device: ", CC.device, f"({torch.cuda.get_device_name(CC.device)})" if torch.cuda.is_available() else "")
        model = CLIP(
            CC.emb_dim, CC.vit_width, CC.img_size, CC.patch_size, CC.n_channels, 
            CC.vit_layers, CC.vit_heads, CC.vocab_size, CC.text_width, 
            CC.max_seq_length, CC.text_heads, CC.text_layers
        ).to(CC.device)
        optimizer = optim.Adam(model.parameters(), lr=CC.lr)
        best_loss = np.inf
        for epoch in range(CC.epochs):
            for i, data in enumerate(train_loader, 0):
                img, cap, mask = data["image"].to(CC.device), data["caption"].to(CC.device), data["mask"].to(CC.device)
                loss = model(img,cap,mask)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            # Saves model if it performed better than the previous best
            print(f"Epoch [{epoch+1}/{CC.epochs}], Batch Loss: {loss.item():.3f}")
            if loss.item() <= best_loss:
                best_loss = loss.item()
                torch.save(model.state_dict(), CC.ckpt_fn)
                print("Model Saved.")

    @staticmethod
    def test_main(params:Dict = {}) -> None:
        # Loading Best Model
        model = CLIP(
            CC.emb_dim, CC.vit_width, CC.img_size, CC.patch_size, CC.n_channels, 
            CC.vit_layers, CC.vit_heads, CC.vocab_size, CC.text_width, 
            CC.max_seq_length, CC.text_heads, CC.text_layers
        ).to(CC.device)
        model.load_state_dict(torch.load(CC.ckpt_fn, map_location=CC.device))
        test_set = FashionMNIST(train = False)
        test_loader = DataLoader(test_set, shuffle=False, batch_size=CC.batch_size)

        # Getting dataset captions to compare images to
        text = torch.stack([tokenizer(x)[0] for x in test_set.captions.values()]).to(CC.device)
        mask = torch.stack([tokenizer(x)[1] for x in test_set.captions.values()])
        mask = mask.repeat(1,len(mask[0])).reshape(len(mask),len(mask[0]),len(mask[0])).to(CC.device)

        correct, total = 0,0
        with torch.no_grad():
            for data in test_loader:
                images, labels = data["image"].to(CC.device), data["caption"].to(CC.device)
                image_features = model.image_encoder(images)
                text_features = model.text_encoder(text, mask=mask)

                image_features /= image_features.norm(dim=-1, keepdim=True)
                text_features /= text_features.norm(dim=-1, keepdim=True)
                similarity = (100.0 * (image_features @ text_features.T)).softmax(dim=-1)
                _, indices = torch.max(similarity,1)
                pred = torch.stack([tokenizer(test_set.captions[int(i)])[0] for i in indices]).to(CC.device)
                correct += int(sum(torch.sum((pred==labels),dim=1)//len(pred[0])))
                total += len(labels)

        print(f'\nModel Accuracy: {100 * correct // total} %')

    @staticmethod
    def zero_shot(params:Dict = {}) -> None:
        # Loading Best Model
        model = CLIP(
            CC.emb_dim, CC.vit_width, CC.img_size, CC.patch_size, CC.n_channels, 
            CC.vit_layers, CC.vit_heads, CC.vocab_size, CC.text_width, CC.max_seq_length, 
            CC.text_heads, CC.text_layers
        ).to(CC.device)
        model.load_state_dict(torch.load(CC.ckpt_fn, map_location=CC.device))
        test_set = FashionMNIST(train = False)
        test_loader = DataLoader(test_set, shuffle=False, batch_size=CC.batch_size)


        # Captions to compare images to
        class_names =["t-shirt/top",
                                "trousers",
                                "pullover",
                                "dress",
                                "coat",
                                "sandal",
                                "shirt",
                                "sneaker",
                                "bag",
                                "ankle boot"]

        text = torch.stack([tokenizer(x)[0] for x in class_names]).to(CC.device)
        mask = torch.stack([tokenizer(x)[1] for x in class_names])
        mask = mask.repeat(1,len(mask[0])).reshape(len(mask),len(mask[0]),len(mask[0])).to(CC.device)

        idx = 1000

        img = test_set[idx]["image"][None,:]
        plt.imshow(img[0].permute(1, 2, 0)  ,cmap="gray")
        plt.title(tokenizer(test_set[idx]["caption"], encode=False, mask=test_set[idx]["mask"][0])[0])
        plt.show()
        img = img.to(CC.device)
        with torch.no_grad():
            image_features = model.image_encoder(img)
            text_features = model.text_encoder(text, mask=mask)


        image_features /= image_features.norm(dim=-1, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
        values, indices = similarity[0].topk(5)

        # Print the result
        print("\nTop predictions:\n")
        for value, index in zip(values, indices):
            print(f"{class_names[int(index)]:>16s}: {100 * value.item():.2f}%")

    @staticmethod
    def tost_main(args:argparse.Namespace = None) -> None:
        TVU.init_distributed_mode(args)
        params = vars(args)
        device = torch.device(args.device)
        # fix the seed for reproducibility
        seed = args.seed + TVU.get_rank()
        torch.manual_seed(seed)
        np.random.seed(seed)
        cudnn.benchmark = True
        # 加载数据集
        dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
        dataset_val, _ = build_dataset(is_train=False, args=args)
        if True:  # args.distributed:
            num_tasks = TVU.get_world_size()
            global_rank = TVU.get_rank()
            if args.repeated_aug:
                sampler_train = RASampler(
                    dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
                )
            else:
                sampler_train = torch.utils.data.DistributedSampler(
                    dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
                )
            if args.dist_eval:
                if len(dataset_val) % num_tasks != 0:
                    print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
                        'This will slightly alter validation results as extra duplicate entries are added to achieve '
                        'equal num of samples per-process.')
                sampler_val = torch.utils.data.DistributedSampler(
                    dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
            else:
                sampler_val = torch.utils.data.SequentialSampler(dataset_val)
        else:
            sampler_train = torch.utils.data.RandomSampler(dataset_train)
            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
        data_loader_train = torch.utils.data.DataLoader(
            dataset_train, sampler=sampler_train,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
            pin_memory=args.pin_mem,
            drop_last=True,
        )
        data_loader_val = torch.utils.data.DataLoader(
            dataset_val, sampler=sampler_val,
            batch_size=int(1.5 * args.batch_size),
            num_workers=args.num_workers,
            pin_memory=args.pin_mem,
            drop_last=False
        )
        # 创建并加载模型
        
        mixup_fn = None
        mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
        if mixup_active:
            mixup_fn = Mixup(
                mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
                prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
                label_smoothing=args.smoothing, num_classes=args.nb_classes)
        print(f"Creating model: {args.model}")
        model = create_model(
            args.model,
            pretrained=False,
            num_classes=args.nb_classes,
            drop_rate=args.drop,
            drop_path_rate=args.drop_path,
            drop_block_rate=None
        )
        if args.pretrained:
            if args.pretrained.startswith('https'):
                checkpoint = torch.hub.load_state_dict_from_url(
                    args.pretrained, map_location='cpu', check_hash=True)
            else:
                checkpoint = torch.load(args.pretrained, map_location='cpu', weights_only=False)
            checkpoint_model = checkpoint['model']
            state_dict = model.state_dict()
            for k in ['head.weight', 'head.bias']:
                if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
                    print(f"Removing key {k} from pretrained checkpoint")
                    del checkpoint_model[k]
            model.load_state_dict(checkpoint_model, strict=True)
        model.to(device)
        if args.surgery:
            checkpoint = torch.load(args.surgery, map_location='cpu')
            checkpoint_model = checkpoint['model']
            patch_embed_weights = {key.replace("patch_embed.", ""): value for key,
                                value in checkpoint['model'].items() if 'patch_embed' in key}
            model.patch_embed.load_state_dict(patch_embed_weights)
            for p in model.patch_embed.parameters():
                p.requires_grad = False
        model_ema = None
        if args.model_ema:
            # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
            model_ema = ModelEma(
                model,
                decay=args.model_ema_decay,
                device='cpu' if args.model_ema_force_cpu else '',
                resume='')
        model_without_ddp = model
        if args.distributed:
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
            model_without_ddp = model.module
        n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print('number of params:', n_parameters)
        linear_scaled_lr = args.lr * args.batch_size * TVU.get_world_size() / 512.0
        args.lr = linear_scaled_lr
        optimizer = create_optimizer(args, model_without_ddp)
        loss_scaler = NativeScaler()
        lr_scheduler, _ = create_scheduler(args, optimizer)
        criterion = LabelSmoothingCrossEntropy()
        if args.mixup > 0.:
            # smoothing is handled with mixup label transform
            criterion = SoftTargetCrossEntropy()
        elif args.smoothing:
            criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
        else:
            criterion = torch.nn.CrossEntropyLoss()
        teacher_model = None
        if args.distillation_type != 'none':
            assert args.teacher_path, 'need to specify teacher-path when using distillation'
            print(f"Creating teacher model: {args.teacher_model}")
            teacher_model = create_model(
                args.teacher_model,
                pretrained=False,
                num_classes=args.nb_classes,
                global_pool='avg',
            )
            if args.teacher_path.startswith('https'):
                checkpoint = torch.hub.load_state_dict_from_url(
                    args.teacher_path, map_location='cpu', check_hash=True)
            else:
                checkpoint = torch.load(args.teacher_path, map_location='cpu')

            teacher_model.load_state_dict(checkpoint['model'])
            teacher_model.to(device)
            teacher_model.eval()
        # wrap the criterion in our custom DistillationLoss, which
        # just dispatches to the original criterion if args.distillation_type is 'none'
        criterion = DistillationLoss(
            criterion, teacher_model, args.distillation_type, args.distillation_alpha, args.distillation_tau
        )
        output_dir = Path(args.output_dir)
        if not os.path.exists(output_dir):
            os.mkdir(output_dir)
        resume_path = os.path.join(output_dir, 'clip.pth')
        if args.resume and os.path.exists(resume_path):
            if args.resume.startswith('https'):
                checkpoint = torch.hub.load_state_dict_from_url(
                    args.resume, map_location='cpu', check_hash=True)
            else:
                print("Loading from checkpoint ...")
                checkpoint = torch.load(resume_path, map_location='cpu')
            model_without_ddp.load_state_dict(checkpoint['model'])
            if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
                optimizer.load_state_dict(checkpoint['optimizer'])
                lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
                args.start_epoch = checkpoint['epoch'] + 1
                if args.model_ema:
                    TVU._load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
                if 'scaler' in checkpoint:
                    loss_scaler.load_state_dict(checkpoint['scaler'])
        if args.eval:
            test_stats = evaluate(data_loader_val, model, device)
            print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
            return
        if args.predict:
            predict(args, dataset_train, sampler_train, model, device)
            return
        # 训练模型
        print(f"Start training for {args.epochs} epochs")
        start_time = time.time()
        max_accuracy = 0.0
        for epoch in range(args.start_epoch, args.epochs):
            if args.distributed:
                data_loader_train.sampler.set_epoch(epoch)
            train_stats = train_one_epoch(
                model, criterion, data_loader_train,
                optimizer, device, epoch, loss_scaler,
                args.clip_grad, model_ema, mixup_fn,
                surgery=args.surgery
            )
            lr_scheduler.step(epoch)
            if args.output_dir:
                checkpoint_paths = [output_dir / 'checkpoint.pth']
                for checkpoint_path in checkpoint_paths:
                    TVU.save_on_master({
                        'model': model_without_ddp.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'lr_scheduler': lr_scheduler.state_dict(),
                        'epoch': epoch,
                        'model_ema': get_state_dict(model_ema),
                        'scaler': loss_scaler.state_dict(),
                        'args': args,
                    }, checkpoint_path)
            if (epoch % args.test_freq == 0) or (epoch == args.epochs - 1):
                test_stats = evaluate(data_loader_val, model, device)

                if test_stats["acc1"] >= max_accuracy:
                    TVU.save_on_master({
                        'model': model_without_ddp.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'lr_scheduler': lr_scheduler.state_dict(),
                        'epoch': epoch,
                        'model_ema': get_state_dict(model_ema),
                        'args': args,
                    }, os.path.join(output_dir, 'best_model.pth'))
                print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
                max_accuracy = max(max_accuracy, test_stats["acc1"])
                print(f'Max accuracy: {max_accuracy:.2f}%')

                log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                            **{f'test_{k}': v for k, v in test_stats.items()},
                            'epoch': epoch,
                            'n_parameters': n_parameters}
                if args.output_dir and TVU.is_main_process():
                    with (output_dir / "log.txt").open("a") as f:
                        f.write(json.dumps(log_stats) + "\n")
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        print('Training time {}'.format(total_time_str))
        print(f'^_^ The End! v0.0.1 ^_^')


    @staticmethod
    def get_args_parser():
        parser = argparse.ArgumentParser('ToST training and evaluation script', add_help=False)
        parser.add_argument('--batch-size', default=64, type=int)
        parser.add_argument('--epochs', default=400, type=int)

        # Model parameters
        parser.add_argument('--model', default='tost_cab_tiny_12_p16', type=str, metavar='MODEL',
                            help='Name of model to train')
        parser.add_argument('--input-size', default=224, type=int, help='images input size')

        parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
                            help='Dropout rate (default: 0.)')
        parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT',
                            help='Drop path rate (default: 0.1)')

        parser.add_argument('--model-ema', action='store_true')
        parser.add_argument('--no-model-ema', action='store_false', dest='model_ema')
        parser.set_defaults(model_ema=True)
        parser.add_argument('--model-ema-decay', type=float, default=0.99996, help='')
        parser.add_argument('--model-ema-force-cpu', action='store_true', default=False, help='')

        # Optimizer parameters
        parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
                            help='Optimizer (default: "adamw"')
        parser.add_argument('--opt-eps', default=1e-8, type=float, metavar='EPSILON',
                            help='Optimizer Epsilon (default: 1e-8)')
        parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA',
                            help='Optimizer Betas (default: None, use opt default)')
        parser.add_argument('--clip-grad', type=float, default=None, metavar='NORM',
                            help='Clip gradient norm (default: None, no clipping)')
        parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                            help='SGD momentum (default: 0.9)')
        parser.add_argument('--weight-decay', type=float, default=0.05,
                            help='weight decay (default: 0.05)')

        # Learning rate schedule parameters
        parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER',
                            help='LR scheduler (default: "cosine"')
        parser.add_argument('--lr', type=float, default=5e-4, metavar='LR',
                            help='learning rate (default: 5e-4)')
        parser.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct',
                            help='learning rate noise on/off epoch percentages')
        parser.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT',
                            help='learning rate noise limit percent (default: 0.67)')
        parser.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV',
                            help='learning rate noise std-dev (default: 1.0)')
        parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR',
                            help='warmup learning rate (default: 1e-6)')
        parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',
                            help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')

        parser.add_argument('--decay-epochs', type=float, default=30, metavar='N',
                            help='epoch interval to decay LR')
        parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',
                            help='epochs to warmup LR, if scheduler supports')
        parser.add_argument('--cooldown-epochs', type=int, default=10, metavar='N',
                            help='epochs to cooldown LR at min_lr, after cyclic schedule ends')
        parser.add_argument('--patience-epochs', type=int, default=10, metavar='N',
                            help='patience epochs for Plateau LR scheduler (default: 10')
        parser.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE',
                            help='LR decay rate (default: 0.1)')

        # Augmentation parameters
        parser.add_argument('--color-jitter', type=float, default=0.4, metavar='PCT',
                            help='Color jitter factor (default: 0.4)')
        parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
                            help='Use AutoAugment policy. "v0" or "original". " + \
                                "(default: rand-m9-mstd0.5-inc1)'),
        parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
        parser.add_argument('--train-interpolation', type=str, default='bicubic',
                            help='Training interpolation (random, bilinear, bicubic default: "bicubic")')

        parser.add_argument('--repeated-aug', action='store_true')
        parser.add_argument('--no-repeated-aug', action='store_false', dest='repeated_aug')
        parser.set_defaults(repeated_aug=True)

        # * Random Erase params
        parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
                            help='Random erase prob (default: 0.25)')
        parser.add_argument('--remode', type=str, default='pixel',
                            help='Random erase mode (default: "pixel")')
        parser.add_argument('--recount', type=int, default=1,
                            help='Random erase count (default: 1)')
        parser.add_argument('--resplit', action='store_true', default=False,
                            help='Do not random erase first (clean) augmentation split')

        # * Mixup params
        parser.add_argument('--mixup', type=float, default=0.8,
                            help='mixup alpha, mixup enabled if > 0. (default: 0.8)')
        parser.add_argument('--cutmix', type=float, default=1.0,
                            help='cutmix alpha, cutmix enabled if > 0. (default: 1.0)')
        parser.add_argument('--cutmix-minmax', type=float, nargs='+', default=None,
                            help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
        parser.add_argument('--mixup-prob', type=float, default=1.0,
                            help='Probability of performing mixup or cutmix when either/both is enabled')
        parser.add_argument('--mixup-switch-prob', type=float, default=0.5,
                            help='Probability of switching to cutmix when both mixup and cutmix enabled')
        parser.add_argument('--mixup-mode', type=str, default='batch',
                            help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')

        # Distillation parameters
        parser.add_argument('--teacher-model', default='regnety_160', type=str, metavar='MODEL',
                            help='Name of teacher model to train (default: "regnety_160"')
        parser.add_argument('--teacher-path', type=str, default='')
        parser.add_argument('--distillation-type', default='none', choices=['none', 'soft', 'hard'], type=str, help="")
        parser.add_argument('--distillation-alpha', default=0.5, type=float, help="")
        parser.add_argument('--distillation-tau', default=1.0, type=float, help="")


        # Dataset parameters
        parser.add_argument('--data-path', default='work/datas/cifar100', type=str,
                            help='dataset path')
        parser.add_argument('--data-set', default='CIFAR100', choices=['CIFAR10', 'CIFAR100', 'IMNET',
                                                                    'INAT', 'INAT19', 'CARS', 'FLOWERS',
                                                                    'IMNET22k'],
                            type=str, help='Image Net dataset path')
        parser.add_argument('--inat-category', default='name',
                            choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
                            type=str, help='semantic granularity')

        parser.add_argument('--output_dir', default='',
                            help='path where to save, empty for no saving')
        parser.add_argument('--device', default='cuda',
                            help='device to use for training / testing')
        parser.add_argument('--seed', default=0, type=int)
        parser.add_argument('--resume', default='', help='resume from checkpoint')
        parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                            help='start epoch')
        parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
        parser.add_argument('--predict', action='store_true', help='Perform evaluation only')
        parser.add_argument('--dist-eval', action='store_true', default=False, help='Enabling distributed evaluation')
        parser.add_argument('--num_workers', default=10, type=int)
        parser.add_argument('--pin-mem', action='store_true',
                            help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
        parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem',
                            help='')
        parser.set_defaults(pin_mem=True)

        # distributed training parameters
        parser.add_argument('--world_size', default=1, type=int,
                            help='number of distributed processes')
        parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
        parser.add_argument('--test-freq', default=1, type=int, help='Number of epochs between \
                                                                    validation runs.')

        parser.add_argument('--full_crop', action='store_true', help='use crop_ratio=1.0 instead of the\
                                                                    default 0.875 (Used by CaiT).')
        parser.add_argument("--pretrained", default=None, type=str, help='Path to pre-trained checkpoint')
        parser.add_argument('--surgery', default=None, type=str, help='Path to checkpoint to copy the \
                                                                    patch projection from. \
                                                                    Can improve stability for very \
                                                                    large models.')

        return parser